package com.yuxl.common.core.domain.socket;

import org.springframework.web.reactive.socket.WebSocketSession;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;

/**
 * @Author : yuxl
 * @create 2025/7/10 下午1:49
 * socket 连接管理，支持多端多连接
 */
public class Connects {
    // userId -> 该用户所有连接
    private static final ConcurrentMap<String, Set<WebSocketSession>> userSessions = new ConcurrentHashMap<>();

    public static void addSession(String userId, WebSocketSession session) {
        userSessions.computeIfAbsent(userId, k -> ConcurrentHashMap.newKeySet()).add(session);
    }

    public static void removeSession(String userId, WebSocketSession session) {
        Set<WebSocketSession> sessions = userSessions.get(userId);
        if (sessions != null) {
            sessions.remove(session);
            if (sessions.isEmpty()) {
                userSessions.remove(userId);
            }
        }
    }

    public static Set<WebSocketSession> getSessions(String userId) {
        return userSessions.getOrDefault(userId, Set.of());
    }

    public static Set<WebSocketSession> all() {
        return userSessions.values().stream().flatMap(Set::stream).collect(Collectors.toSet());
    }
}
